# -*- coding: utf-8 -*-
"""
@Time ： 2020-11-30 16:19
@Auth ： liangpw3
@Description：

"""

from algo.Algo_interface import Algo_interface
# import api as runs
import json
from sklearn.decomposition import PCA
from sklearn.neural_network import MLPClassifier
import numpy as np


class ImageTypeClassifier(Algo_interface):
    def __init__(self, model_type, model_name, model_params):
        self.task_type = model_type
        self.model_name = model_name
        self.model_params = model_params
        self.model = dict()
        self.build_model()
        # return self.model

    def set_model(self, model):
        self.model = model
        return 1

    def get_model(self):
        return self.model

    def build_model(self):
        if self.model_name == '图像分类-v1':
            # n_components后期改超参,都可以直接由model_params拆解出来，如果换模型另算
            self.model['pca'] = PCA(n_components=128, whiten=True)
            self.model['model'] = MLPClassifier(**self.model_params)

    def train(self,data):
        x_train, y_train = data  #在api.py是tuple进来的两个Series
        x_train = np.asarray(x_train.apply(lambda x: x.flatten()).tolist())    #对df内的图像array光栅化并转矩阵
        x_train_pca = self.model['pca'].fit_transform(x_train)   # fit pca
        self.model['model'].fit(x_train_pca, y_train)
        return 1

    def predict(self,data):   #predict输入就是x_test了，Series格式
        x_test = np.asarray(data.apply(lambda x: x.flatten()).tolist())   #对df内的图像array光栅化并转矩阵
        x_test_pca = self.model['pca'].transform(x_test)   #
        y_pred = self.model['model'].predict(x_test_pca)
        return y_pred